package ffte.core.flat

import spinal.core._
import spinal.lib._

import ffte.types.VecComplex
import ffte.types.FixComplex
import ffte.property.{getFFTIW,getFFTOW,FFTIOConfig,getTwiddleResolution}
import ffte.algorithm.{circularOrder,FFTMisc}
import ffte.algorithm.FFTGen.{Method,wpbFFT}

/**
 * Winograd Prime-Based (WPB) FFT implementation
 *
 * This class implements FFT using the Winograd Prime-Based decomposition method,
 * which combines Winograd's minimal multiplication algorithms with prime factor
 * decomposition. WPB is optimal for transform sizes that are products of small
 * primes, offering the minimal number of multiplications.
 *
 * Mathematical Foundation:
 * The WPB method uses Winograd's theorem which states that for small prime
 * lengths (2, 3, 4, 5, 7, 8, 9, 16), the DFT can be computed with fewer
 * multiplications than the standard Cooley-Tukey approach. When combined
 * with prime factor decomposition for N = p1^e1 * p2^e2 * ... * pk^ek,
 * this yields optimal algorithms.
 *
 * Winograd Algorithm Characteristics:
 * - Minimal multiplication count for small prime lengths
 * - Increased number of additions/subtractions
 * - Pre/post-processing with addition matrices
 * - Pointwise multiplication in transformed domain
 *
 * Hardware Architecture:
 * - Special preprocessing: Input permutation and addition network
 * - Core computation: Reduced-size FFT on transformed inputs
 * - Twiddle multiplication: Optimized for Winograd structure
 * - Postprocessing: Output addition network and permutation
 *
 * Performance Characteristics:
 * - Minimal multiplication count for prime-based sizes
 * - Higher addition count compared to Cooley-Tukey
 * - Excellent for ASIC implementation where multipliers are expensive
 * - Regular structure suitable for hardware implementation
 *
 * @tparam T Method type for sub-FFT implementations (Flat, Mixed, Streamed)
 * @param S Scaling shift amount for fixed-point overflow prevention
 * @param p WPB decomposition structure containing Winograd algorithm information
 */
class wpb[T <: Method](val S: Int, val p: wpbFFT[T]) extends combineFlatComponents[T] {
    /** Standard I/O interface for flat FFT components */
    val io = iogen

    /**
     * Twiddle factor precision configuration
     *
     * The twiddle factors use reduced precision (1 - resolution) to optimize
     * hardware resource usage while maintaining sufficient accuracy for
     * the Winograd algorithm requirements.
     */
    val tW = 1 - getTwiddleResolution()

    /**
     * Circular ordering configuration for Winograd algorithm
     *
     * The Winograd algorithm requires specific input/output ordering
     * for optimal implementation. The circular order provides the
     * permutation pattern needed for the pre/post-processing stages.
     */
    val oc = circularOrder(p.N)

    /**
     * Precision increase for intermediate calculations
     *
     * Winograd algorithms require additional precision to handle
     * the increased dynamic range from the addition matrices.
     * This is calculated based on the transform size and scaling.
     */
    val increase = Math.round(Math.log(p.N - 1) / Math.log(2) - S).toInt

    /**
     * Winograd twiddle factors and precision requirements
     *
     * The circular ordering object provides:
     * - wf: Twiddle factors with optimal fixed-point representation
     * - ws: Additional precision shift for output processing
     */
    val (wf, ws) = oc.wfft(tW)

    /**
     * Output buffer with increased precision
     *
     * Winograd algorithms produce intermediate results with increased
     * dynamic range, requiring additional precision for accurate
     * computation before final scaling.
     */
    val o = VecComplex(p.N, oW + increase)

    /**
     * First stage: Core FFT computation
     *
     * This block implements the main FFT computation using the Winograd
     * approach. It processes the permuted inputs through a reduced-size
     * FFT that takes advantage of the Winograd minimal multiplication
     * algorithms.
     *
     * Architecture:
     * - Single FFT processor of size (p.N - 1) for non-DC components
     * - Input width: dW (global input precision)
     * - Output width: interW (intermediate precision)
     * - Uses flat implementation for combinatorial computation
     */
    val f_block = FFTIOConfig(dW, interW) on new Area {
        /**
         * Winograd FFT processor for non-DC components
         *
         * Creates a single FFT processor that handles all non-DC
         * components using the Winograd algorithm structure.
         */
        val u = fg(first, s"f")
        val delay = u.delay
        u.io.ce := io.ce
    }

    /**
     * Second stage: Additional processing
     *
     * This block implements the second stage of the Winograd decomposition,
     * handling any remaining processing required for the complete algorithm.
     *
     * Architecture:
     * - Flat FFT processor with second stage scaling
     * - Input width: interW (intermediate precision)
     * - Output width: oW + increase (increased precision for final processing)
     */
    val s_block = FFTIOConfig(interW, oW + increase) on new Area {
        /**
         * Second stage Winograd processor
         *
         * Uses the Flat method with appropriate scaling for the
         * second stage of the Winograd decomposition.
         */
        val g = new Flat(second_shift)
        val u = g(second, s"f")
        val delay = u.delay
        u.io.ce := io.ce
    }

    /**
     * Total delay through processing stages
     *
     * Includes 1 cycle for initial register stages plus delays
     * through both FFT processing blocks.
     */
    val fdelay = 1 + f_block.delay + s_block.delay

    /**
     * Intermediate result buffer register
     *
     * Stores the twiddle-multiplied results from the first stage
     * before feeding them to the second stage. This register breaks
     * the combinatorial path for better timing.
     *
     * Size: p.N - 1 (excluding DC component)
     * Width: interW bits (intermediate precision)
     */
    val xf = Reg(VecComplex(p.N - 1, interW))

    /**
     * Connect buffer to second stage input
     *
     * The buffered intermediate results are fed directly to the
     * second stage FFT processor for further computation.
     */
    s_block.u.io.d := xf

    /**
     * Winograd processing for non-DC components
     *
     * This loop implements the core Winograd algorithm processing
     * for all non-DC components (indices 1 to p.N-1). It includes
     * twiddle factor multiplication and proper precision management.
     *
     * Processing Steps:
     * 1. Input routing (excluding DC component)
     * 2. Twiddle factor generation and multiplication
     * 3. Precision management through cutting
     * 4. Output precision adjustment
     */
    for (i <- 1 until p.N) {
        /**
         * Input to first stage connection
         *
         * Routes input data (excluding DC component at index 0) to
         * the first stage Winograd FFT processor.
         */
        f_block.u.io.d.d(i - 1) := io.d.d(i)

        /**
         * Winograd twiddle factor generation
         *
         * Creates a complex twiddle factor with the precision required
         * for Winograd algorithms. The factor is generated using the
         * circular ordering table for optimal permutation.
         */
        val f = FixComplex(1 - getTwiddleResolution(), getTwiddleResolution())
        f.re.fromInt(wf(p.p.first.store_tab(i - 1))._1)
        f.im.fromInt(wf(p.p.first.store_tab(i - 1))._2)

        /**
         * Twiddle multiplication and precision management
         *
         * Multiplies the first stage output by the Winograd twiddle factor
         * and manages precision through cutting operations to prevent
         * overflow while maintaining accuracy.
         */
        when(io.ce) {
            xf.d(i - 1) := (f * f_block.u.io.q.d(i - 1)).cut(interW, -getTwiddleResolution(), f_block.u.io.q.d(i - 1).resolution)
        }

        /**
         * Output precision adjustment
         *
         * Adjusts the second stage output precision to match the
         * increased precision required by the Winograd algorithm.
         */
        o.d(i) := s_block.u.io.q.d(i - 1).resize(oW + increase + ws).cut(oW + increase, ws)
    }

    /**
     * DC component processing
     *
     * The DC component (index 0) is processed separately using a
     * specialized sum operation that accumulates all non-DC inputs.
     * This implements the special handling required by Winograd
     * algorithms for the zero-frequency component.
     */
    o.d(0) := FixComplex.sumc(Vec((1 until p.N).map(io.d.d(_).resize(oW + increase + S))), p.N - 1, 0, fdelay, io.ce).cut(oW + increase, S)

    /**
     * Delayed DC input for final addition
     *
     * The original DC input is delayed to match the processing
     * latency of the non-DC components, enabling proper final
     * combination of all components.
     */
    val d0 = Delay(io.d.d(0), fdelay, io.ce).resize(oW + increase + S).cut(oW + increase, S)

    /**
     * Final output buffer registers
     *
     * Stores the final FFT results after combining the processed
     * DC component with the Winograd-processed non-DC components.
     * Uses registers for proper timing control.
     *
     * Size: p.N (all frequency components)
     * Width: oW bits (final output precision)
     */
    val q = Reg(VecComplex(p.N, oW))

    /**
     * Final output combination and scaling
     *
     * Combines the delayed DC input with the Winograd-processed
     * components and scales to the final output precision.
     * This completes the Winograd FFT computation.
     */
    when(io.ce) {
        for (i <- 0 until p.N) q.d(i) := (d0 + o.d(i)).fixTo(oW, 0)
    }

    /**
     * Output connection
     *
     * Routes the registered final results to the component output.
     */
    io.q := q

    /**
     * Total component delay in clock cycles
     *
     * Total delay includes all processing stages plus one additional
     * cycle for the final output registers.
     *
     * @return Total delay from input to output in clock cycles
     */
    def delay = fdelay + 1
}
